#!/usr/bin/env python3

#==============================================================================
# Functionality
#==============================================================================
import pdb
import sys
import os
import re

# utility funcs, classes, etc go here.

def asserting(cond):
    if not cond:
        pdb.set_trace()
    assert(cond)

def has_stdin():
    return not sys.stdin.isatty()

def reg(pat, flags=0):
    return re.compile(pat, re.VERBOSE | flags)

#==============================================================================
# Cmdline
#==============================================================================
import argparse

parser = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter, 
    description="""
TODO
""")
     
parser.add_argument('-v', '--verbose',
    action="store_true",
    help="verbose output" )
     
parser.add_argument('-c', '--compact',
    action="store_true",
    help="one JSON entry per line" )
     
parser.add_argument('-u', '--upload-images',
    action="store_true",
    help="try to replace base64-encoded images with a permanent url" )

parser.add_argument('-w', '--walltime-min',
    type=float,
    default=None,
    help="Skip summaies whose walltime is < this value" )

parser.add_argument('-s', '--step-min',
    type=int,
    default=None,
    help="Skip summaies whose step is < this value" )

args = None

#==============================================================================
# Main
#==============================================================================

from base64 import b64decode
import re
from tensorboard.compat.proto import event_pb2
import google.protobuf.json_format as jf
import json
import traceback

def replace_upload(m):
  print(m);

def process_base64_events(f):
  for line in f:
    if ',' in line:
      line = line.rsplit(',', 1)[-1]
    event_str = b64decode(line)
    #print(len(event_str))
    event = event_pb2.Event.FromString(event_str)
    if args.step_min is not None or args.walltime_min is not None:
      if args.step_min is not None and event.step < args.step_min:
        continue
      if args.walltime_min is not None and event.wall_time < args.walltime_min:
        continue
    #print(event)
    if args.compact:
      js = jf.MessageToJson(event, indent=0)
      js = ''.join(js.splitlines())
    else:
      js = jf.MessageToJson(event)
      # if args.compact:
      #   js = json.dumps(json.loads(js))
    #if args.upload_images:
    #  js = re.sub(reg(r""" "encodedImageString" \s* : \s* "(.*?)" """), replace_upload)
    print(js)
    # try:
    # except:
    #   traceback.print_exc()
    #   sys.stderr.write('For event string %s\n' % repr(event_str))
    #   sys.stderr.flush()



def run():
    if args.verbose:
        print(args)
    if len(args.args) <= 0 and not has_stdin():
        # if there were no args and there was no input, prompt user.
        print('Enter input (press Ctrl-D when done):')
    if len(args.args) <= 0 or has_stdin():
        process_base64_events(sys.stdin)
        return
    # for each arg on cmdline...
    for arg in args.args:
        print(arg)

def main():
    try:
        global args
        if not args:
            args, leftovers = parser.parse_known_args()
            args.args = leftovers
        return run()
    except IOError:
        # http://stackoverflow.com/questions/15793886/how-to-avoid-a-broken-pipe-error-when-printing-a-large-amount-of-formatted-data
        try:
            sys.stdout.close()
        except IOError:
            pass
        try:
            sys.stderr.close()
        except IOError:
            pass

if __name__ == "__main__":
    main()

